import os
import requests
import time
import re
from concurrent.futures import ThreadPoolExecutor


def download_image(image_url, file_path):
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/121.0.0.0 Safari/537.36'
    }

    attempts = 5  # 最大重试次数
    for attempt in range(attempts):
        try:
            response = requests.get(image_url, headers=headers)
            response.raise_for_status()  # 检查请求是否成功

            # 保存为 PNG 格式（只是文件扩展名为 .png，内容依然为 webp）
            with open(file_path, 'wb') as file:
                file.write(response.content)

            print(f"Downloaded (saved as PNG): {file_path}")
            return  # 成功下载，退出函数

        except Exception as e:
            print(f"Failed to download {image_url}: {e}")
            time.sleep(5)  # 等待 1 秒后重试

    print(f"Max attempts reached. Skipping {image_url}.")


def download_images(first_image_url, count):
    # 从完整的图片 URL 中提取基础 URL 和文件名
    base_url = first_image_url.rsplit('/', 1)[0] + '/'  # 基础 URL
    initial_id = first_image_url.split('/')[-1]  # 获取文件名

    # 查找"_result"的位置
    result_index = initial_id.find('_result')
    if result_index == -1:
        print("文件名中未找到 _result。")
        return  # 如果没有找到，退出程序

    # 提取基本名称及其编号
    base_name = initial_id[:result_index]  # 基本名称
    match = re.search(r'_(\d+)$', base_name)  # 使用正则表达式提取最后的数字部分

    if match:
        start_number = int(match.group(1))  # 提取数字并转换为整数
        name_part = base_name[:match.start()]  # 提取 "_result" 前面的部分
    else:
        print("未能从文件名中提取数字，退出程序。")
        return  # 如果没有找到数字，退出程序

    # 创建文件夹
    folder_path = name_part  # 使用基本名称作为文件夹名称
    os.makedirs(folder_path, exist_ok=True)  # 创建目录，如果已存在则不抛出错误

    # 创建一个线程池，最多同时运行 5 个线程
    with ThreadPoolExecutor(max_workers=2) as executor:
        futures = []
        for i in range(count):  # 从提取的数字开始编号
            # 生成新的图片 URL，确保编号前面带有下划线
            image_url = f"{base_url}{name_part}_{start_number + i}_result.webp"  # 正确的拼接格式
            print(f"Attempting to download: {image_url}")  # 打印生成的 URL

            # 文件路径
            file_path = os.path.join(folder_path, f"{name_part}_{start_number + i}_result.png")  # 保存为 PNG 格式

            # 检查文件是否已经存在
            if os.path.exists(file_path):
                print(f"Image {file_path} already exists. Skipping.")
                continue

                # 提交下载任务
            futures.append(executor.submit(download_image, image_url, file_path))

            # 等待全部线程完成
        for future in futures:
            future.result()  # 获取每个线程的结果（或抛出异常）


if __name__ == "__main__":
    # 输入完整的第一张图片的 URL 和下载数量
    first_image_url = input("请输入第一张图片的完整 URL: ")
    num_images = int(input("请输入要下载的图片数量: "))

    download_images(first_image_url, num_images)